package com.ds.infra.rpc.registry.service.impl;

import com.ds.infra.rpc.registry.constant.ZkConstant;
import com.ds.infra.rpc.registry.model.ServiceInstance;
import com.ds.infra.rpc.registry.service.ServiceDiscovery;
import com.ds.infra.rpc.registry.util.CuratorUtils;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * 作用：RPC服务发现核心接口
 *
 * @author WeiShaoying
 */
@Slf4j
public class ZkServiceDiscovery implements ServiceDiscovery {

    private final CuratorFramework zkClient;

    private final ObjectMapper objectMapper = new ObjectMapper();

    public ZkServiceDiscovery(String zkAddress) {
        this.zkClient = CuratorUtils.getZkClient(zkAddress);
    }

    @Override
    public List<Map<String, Object>> discover(String serviceKey) throws Exception {
        String servicePath = ZkConstant.ZK_BASE_PATH + "/" + serviceKey;
        log.info("==== Discovering servicePath: {}", servicePath);
        List<String> children = zkClient.getChildren().forPath(servicePath);
        log.info("==== Discovered service nodes: {}", children);

        return children.stream()
                .map(child -> {
                    try {
                        byte[] data = zkClient.getData().forPath(servicePath + "/" + child);
                        return objectMapper.readValue(data, new TypeReference<Map<String, Object>>() {
                        });
                    } catch (Exception e) {
                        log.error("Failed to parse service node: {}", child, e);
                        return null;
                    }
                })
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }

    @Override
    public List<ServiceInstance> discoverInstance(String serviceKey) throws Exception {
        return List.of();
    }

    @Override
    public void watchService(String serviceKey, Runnable changeHandler) throws Exception {
        String servicePath = ZkConstant.ZK_BASE_PATH + "/" + serviceKey;
        PathChildrenCache watcher = new PathChildrenCache(zkClient, servicePath, true);

        watcher.getListenable().addListener((client, event) -> {
            if (event.getType() == PathChildrenCacheEvent.Type.CHILD_ADDED ||
                    event.getType() == PathChildrenCacheEvent.Type.CHILD_REMOVED) {
                changeHandler.run();
            }
        });
        watcher.start(PathChildrenCache.StartMode.BUILD_INITIAL_CACHE);
    }
}
